﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using Roslyn.Compilers.CSharp;
using Roslyn.Services;
using Util;
using CompleterLib.Util;
using CompleterLib;
using System.Text;

namespace QueryGenerator
{
    class QueryableCollector : Roslyn.Compilers.CSharp.SyntaxWalker
    {
        private readonly IDocument doc;
        private readonly SemanticModel semModel;

        private readonly IList<Tuple<InvocationExpressionSyntax, QueryPositionFeatures>> calls;

        public QueryableCollector(IDocument doc)
        {
            this.doc = doc;
            this.calls = new List<Tuple<InvocationExpressionSyntax, QueryPositionFeatures>>();
            var tree = doc.GetSyntaxTree() as SyntaxTree;
            this.semModel = doc.GetSemanticModel() as SemanticModel;

            if (semModel != null && tree != null)
            {
                var ctoken = new CancellationToken();
                this.Visit(tree.GetRoot(ctoken));
            }
        }

        private int localDecls;
        private MethodSymbol currentMethod;
        private PropertySymbol currentProperty;
        private MethodFeatures currentMethodFeatures;

        public override void VisitMethodDeclaration(MethodDeclarationSyntax node)
        {
            localDecls = 0;
            currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
            currentMethodFeatures = new MethodFeatures(currentMethod, semModel);
            base.VisitMethodDeclaration(node);
        }
        public override void VisitPropertyDeclaration(PropertyDeclarationSyntax node)
        {
            localDecls = 0;
            currentMethod = null;
            currentMethodFeatures = null;
            currentProperty = semModel.GetDeclaredSymbol(node) as PropertySymbol;
            base.VisitPropertyDeclaration(node);
        }
        public override void VisitAccessorDeclaration(AccessorDeclarationSyntax node)
        {
            localDecls = 0;
            if (node.Kind.Equals(SyntaxKind.AddAccessorDeclaration) || node.Kind.Equals(SyntaxKind.RemoveAccessorDeclaration))
            {
                currentMethod = null;
                currentMethodFeatures = null;
            }
            else
            {
                try
                {
                    currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
                    currentMethodFeatures = new MethodFeatures(currentMethod, semModel);
                }
                catch (NullReferenceException)
                {
                    currentMethod = null;
                    Logging.Logger.PrintLine("VisitAccessorDeclaration: NullReferenceException"
                        + " in document " + semModel.Compilation.SyntaxTrees.First().FilePath
                        + " for node: " + node);
                }
            }

            base.VisitAccessorDeclaration(node);
        }
        public override void VisitConstructorDeclaration(ConstructorDeclarationSyntax node)
        {
            localDecls = 0;
            currentMethod = semModel.GetDeclaredSymbol(node) as MethodSymbol;
            currentMethodFeatures = new MethodFeatures(currentMethod, semModel);
            base.VisitConstructorDeclaration(node);
        }
        public override void VisitFieldDeclaration(FieldDeclarationSyntax node)
        {
            localDecls = 0;
            currentMethod = null;
            currentMethodFeatures = null;
            base.VisitFieldDeclaration(node);
        }

        public override void VisitLocalDeclarationStatement(LocalDeclarationStatementSyntax node)
        {
            base.VisitLocalDeclarationStatement(node);
            localDecls++;
        }

        public override void VisitInvocationExpression(InvocationExpressionSyntax call)
        {
            base.VisitInvocationExpression(call);

            var method = call.Expression;
            var methodType = semModel.GetTypeInfo(method).Type;
            var methodSymbol = semModel.GetSymbolInfo(method).Symbol;

            if (!(methodType is ErrorTypeSymbol) && methodSymbol != null)
            //&& methodType.TypeKind.Equals(TypeKind.Delegate))
            {
                var meth = methodSymbol as MethodSymbol;
                if (meth == null)
                {
                    return; // XXX ?
                }
                var args = call.ArgumentList.Arguments.Select(a => a.Expression);
                var methodAccess = method as MemberAccessExpressionSyntax;
                if (methodAccess != null)
                {
                    var parent = methodAccess.Expression;
                    if (!(parent is PredefinedTypeSyntax))
                    {
                        // TODO Should parent be included here?
                        var parentInfo = semModel.GetSymbolInfo(parent);
                        if (!(parentInfo.Symbol.ResolveAlias() is TypeSymbol))
                        {
                            args = parent.Cons(args);
                        }
                    }
                }
                else if (!meth.IsStatic)
                {
                    // null means this
                    args = EnumerableUtil.Cons(null, args);
                }
                if (args.Count() >= 2)
                {
                    remember(call);
                }
            }
        }

        private void remember(InvocationExpressionSyntax call)
        {
            calls.Add(new Tuple<InvocationExpressionSyntax, QueryPositionFeatures>(
                call, new QueryPositionFeatures(call, localDecls, doc,
                    semModel, currentMethodFeatures)));
        }

        public IEnumerable<Tuple<InvocationExpressionSyntax, QueryPositionFeatures>> GetCalls()
        {
            return calls;
        }

        public static IEnumerable<MethodQuery<QueryPositionFeatures>> GetQueriesForCall(
            InvocationExpressionSyntax call,
            QueryPositionFeatures queryFeatures, SemanticModel semModel)
        {
            var method = call.Expression;
            var methodType = semModel.GetTypeInfo(method).Type;
            var methodSymbol = semModel.GetSymbolInfo(method).Symbol;

            if (!(methodType is ErrorTypeSymbol) && methodSymbol != null)
            //&& methodType.TypeKind.Equals(TypeKind.Delegate))
            {
                var meth = methodSymbol as MethodSymbol;
                if (meth == null)
                {
                    yield break; // XXX ?
                }
                var args = call.ArgumentList.Arguments.Select(a => a.Expression);
                var methodAccess = method as MemberAccessExpressionSyntax;
                if (methodAccess != null)
                {
                    var parent = methodAccess.Expression;
                    if (!(parent is PredefinedTypeSyntax))
                    {
                        // TODO Should parent be included here?
                        var parentInfo = semModel.GetSymbolInfo(parent);
                        if (!(parentInfo.Symbol.ResolveAlias() is TypeSymbol))
                        {
                            args = parent.Cons(args);
                        }
                    }
                }
                else if (!meth.IsStatic)
                {
                    // null means this
                    // TODO Handle this null somehow.
                    args = EnumerableUtil.Cons(null, args);
                }
                var argSubsets = args
                                .SelectMany(a =>
                                    args.Select(b => new HashSet<ExpressionSyntax>() {a,b}))
                                .Distinct(HashSet<ExpressionSyntax>.CreateSetComparer());
                if (!meth.IsStatic)
                {
                    argSubsets = argSubsets.Where(s => s.Contains(args.First()));
                }
                foreach (var argSubset in argSubsets)
                {
                    yield return new MethodQuery<QueryPositionFeatures>(
                        call, meth, argSubset.ToList(), semModel, queryFeatures);
                }
            }
        }
        public IEnumerable<MethodQuery<QueryPositionFeatures>> GetAllMethodQueries()
        {
            return calls.SelectMany(call =>
                GetQueriesForCall(call.Item1, call.Item2, semModel));
        }
        public IEnumerable<QueryResult<QueryPositionFeatures>> GetAllMethodQueryResults()
        {
            return Completer.PerformMethodQueries(GetAllMethodQueries(), semModel);
        }
        public string PrettyPrintQueryResult(QueryResult<QueryPositionFeatures> result,
            TypeFeatures typeFeatures, MethodFeatures methodFeatures,
            int qid, int maxResults = 100)
        {
            StringBuilder str = new StringBuilder();
            var query = result.Query;
            QueryPositionFeatures queryFeatures = query.ExtraInfo;

            int position = query.OriginalMethodCall.Span.Start;
            var line = doc.GetText().GetLineFromPosition(position);
            int col = position - line.Start;
            int lineNum = line.LineNumber;
            string docname = doc.Name;
            //string posStr = docname + "::" + lineNum + ":" + col;
            str.Append("# Query at ");
            str.Append(docname);
            str.Append("::");
            str.Append(lineNum);
            str.Append(':');
            str.Append(col);
            str.Append(": ");
            str.AppendLine(query.Arguments
                .Select(a => a == null ? "/*implicit*/this" : a.GetText())
                .CommaSeparated());

            foreach(var res in result.Results.Take(maxResults)) {
                bool isCorrect = query.IsCorrectAnswer(res);
                str.Append(string.Join(" ", new object[] {
                    isCorrect.ToIntString(),
                    "qid:" + qid,
                }));
                str.Append(' ');
                IEnumerable<int> features;
                features = res.ScoreInfo.AsEnumerable()
                    .Concat(new int[] {
                        res.Method.Name.Length,
                        res.Method.Name.Count(c => Char.IsUpper(c)),
                        res.Method.ShortName.Length,
                        res.Method.ShortName.Count(c => Char.IsUpper(c)),
                        res.Method.ParameterTypes.Count,
                        res.Method.IsStatic.ToInt(),
                        res.Method.ShortName.Equals("new").ToInt(),
                    })
                    .Concat(typeFeatures.AsEnumerable())
                    .Concat(methodFeatures.AsEnumerable())
                    .Concat(queryFeatures.AsEnumerable());
                //features = res.Score.Singleton();
                str.Append(string.Join(" ",
                    features.Select((f, i) => (i+1) + ":" + f).ToArray()));
                str.Append(" # ");
                str.AppendLine(res.ToString());
            }
            return str.ToString();
        }

        private static int qid = 0;
        internal string PrettyPrintAllMethodQueryResults()
        {
            StringBuilder str = new StringBuilder();
            var typeFeatureCache = new Cache<TypeSymbol, TypeFeatures>(type => new TypeFeatures(type, semModel));
            foreach (var res in GetAllMethodQueryResults())
            {
                var queryFeatures = res.Query.ExtraInfo;
                var methodFeatures = queryFeatures.MethodFeatures;
                var typeFeatures = typeFeatureCache[queryFeatures.ContainingType];
                str.AppendLine(PrettyPrintQueryResult(res, typeFeatures, methodFeatures, qid++));
            }
            return str.ToString();
        }
    }

    public class TypeFeatures
    {
        private readonly TypeSymbol type;

        private readonly Accessibility accessibility;
        private readonly bool isStatic;
        private readonly bool isInner;
        private readonly int nsDepth;
        private readonly int typeDistToObject;

        private readonly int[] allFeatures;

        public TypeFeatures(TypeSymbol type, SemanticModel semModel)
        {
            this.type = type;

            isStatic = type.IsStatic;
            accessibility = type.DeclaredAccessibility;
            isInner = type.ContainingType != null;
            nsDepth = type.ContainingNamespace.Name.Count(c => c.Equals('.'));
            typeDistToObject = semModel.Compilation.ObjectType.DistanceFrom(type, semModel);

            allFeatures = new int[] {
                (int) accessibility,
                accessibility.Equals(Accessibility.Public).ToInt(),
                isStatic.ToInt(),
                isInner.ToInt(),
                nsDepth,
                typeDistToObject,
            };
        }

        internal IEnumerable<int> AsEnumerable()
        {
            return allFeatures;
        }
    }
    public class MethodFeatures
    {
        private readonly MethodSymbol method;

        private readonly Accessibility accessibility;
        private readonly bool isStatic;
        private readonly int numParams;
        private readonly int arity;
        private readonly MethodKind methodKind; // simplified version of this?
        private readonly bool returnsVoid;
        private readonly bool isOverride;
        private readonly int returnTypeDistToObject;

        private readonly int[] allFeatures;

        public MethodFeatures(MethodSymbol sym, SemanticModel semModel)
        {
            this.method = sym;

            if (sym == null)
            {
                accessibility = default(Accessibility);
                numParams = 0;
                arity = 0;
                methodKind = MethodKind.Conversion; // dummy that is not checked for below
                isStatic = true;
                isOverride = false;
                returnsVoid = true;
                returnTypeDistToObject = 0;
            }
            else
            {
                accessibility = sym.DeclaredAccessibility;
                isStatic = sym.IsStatic;
                methodKind = sym.MethodKind;
                numParams = sym.Parameters.Count;
                returnsVoid = sym.ReturnsVoid;
                arity = sym.Arity;
                isOverride = sym.IsOverride;
                returnTypeDistToObject = returnsVoid || !sym.ReturnType.IsValidType()
                    ? 0
                    : semModel.Compilation.ObjectType.DistanceFrom(sym.ReturnType, semModel);
            }

            allFeatures = new int[] {
                (int) accessibility,
                accessibility.Equals(Accessibility.Public).ToInt(),
                isStatic.ToInt(),
                isOverride.ToInt(),
                numParams,
                arity,
                methodKind.Equals(MethodKind.Constructor).ToInt(),
                methodKind.Equals(MethodKind.StaticConstructor).ToInt(),
                methodKind.Equals(MethodKind.PropertyGet).ToInt(),
                methodKind.Equals(MethodKind.PropertySet).ToInt(),
                methodKind.Equals(MethodKind.Ordinary).ToInt(),
                returnsVoid.ToInt(),
                returnTypeDistToObject,
            };
        }

        public MethodSymbol MethodSymbol { get { return method; } }

        internal IEnumerable<int> AsEnumerable()
        {
            return allFeatures;
        }

        private static Lazy<MethodFeatures> nullMethodFeatures =
            new Lazy<MethodFeatures>(() => new MethodFeatures(null, null));
        public static MethodFeatures NullMethodFeatures { get { return nullMethodFeatures.Value; } }
    }
    public enum ExpressionContextKind
    {
        STATEMENT, ARGUMENT, VAR_LOCAL_DECL, LOCAL_DECL, ASSIGNMENT,
        BINARY_ARG, LOGICAL_ARG,
        IF_CONDITION, WHILE_CONDITION,
        UNKNOWN,
        FIELD_DECL,
        LAMBDA,
        RETURN,
        LOOKUP,
        FOR_EACH,
        INITIALIZER
    }
    public class QueryPositionFeatures
    {
        private readonly ExpressionSyntax exp;
        private readonly int numLocalDecls;
        private readonly int expressionDepth;
        private readonly int nestingDepth;
        private readonly int expLineInMethod;
        private readonly ExpressionContextKind contextKind;
        private readonly MethodFeatures methodFeatures;
        private readonly TypeSymbol containingType;

        private readonly int[] allFeatures;

        public QueryPositionFeatures(ExpressionSyntax exp, int numLocalDecls,
            IDocument doc, SemanticModel semModel, MethodFeatures methodFeatures)
        {
            this.exp = exp;
            this.numLocalDecls = numLocalDecls;

            if (methodFeatures == null)
            {
                containingType = semModel.GetContainingTypeAt(exp.Span.Start);
                this.methodFeatures = MethodFeatures.NullMethodFeatures;
            }
            else
            {
                containingType = methodFeatures.MethodSymbol.ContainingType;
                this.methodFeatures = methodFeatures;
            }

            expressionDepth = exp.Ancestors().TakeWhile(n => n is ExpressionSyntax).Count();
            // TODO What if not in a method?
            var containingDeclBlock
                = exp.Ancestors().First(n => (n is BlockSyntax
                        || n is FieldDeclarationSyntax
                        || n is ConstructorDeclarationSyntax)
                    && !(n.Parent is ExpressionSyntax || n.Parent is StatementSyntax));
            nestingDepth = exp.Ancestors().Where(n => n is BlockSyntax).Count();
            var docText = doc.GetText();

            int methodBodyStart = containingDeclBlock.Span.Start;
            int methodBodyStartLine = docText.GetLineNumberFromPosition(methodBodyStart);
            int expStart = exp.Span.Start;
            int expStartLine = docText.GetLineNumberFromPosition(expStart);
            expLineInMethod = expStartLine - methodBodyStartLine;

            contextKind = GetContextKind(exp.Parent, semModel);
            // TODO More features?

            allFeatures = new int[] {
                numLocalDecls,
                expressionDepth,
                nestingDepth,
                expLineInMethod,
            }.Concat(Enum.GetValues(typeof(ExpressionContextKind))
                .Cast<ExpressionContextKind>().Select(e => e.Equals(contextKind).ToInt()))
            .ToArray();
        }

        public MethodFeatures MethodFeatures { get { return methodFeatures; } }

        private static ExpressionContextKind GetContextKind(SyntaxNode parent, SemanticModel semModel)
        {
            if (parent is ExpressionStatementSyntax)
            {
                return ExpressionContextKind.STATEMENT;
            }
            else if (parent is ArgumentListSyntax || parent is ArgumentSyntax)
            {
                return ExpressionContextKind.ARGUMENT;
            }
            else if (parent.Kind.Equals(SyntaxKind.AssignExpression))
            {
                return ExpressionContextKind.ASSIGNMENT;
            }
            else if (parent is IfStatementSyntax || parent is WhileStatementSyntax)
            {
                return ExpressionContextKind.IF_CONDITION;
            }
            else if (parent is ForEachStatementSyntax)
            {
                return ExpressionContextKind.FOR_EACH;
            }
            else if (parent.Kind.Equals(SyntaxKind.AsExpression)
                || parent is EqualsValueClauseSyntax
                || parent is ConditionalExpressionSyntax
                || parent is ParenthesizedExpressionSyntax
                || parent is CastExpressionSyntax)
            {
                return GetContextKind(parent.Parent, semModel);
            }
            else if (parent is VariableDeclaratorSyntax)
            {
                var decl = parent as VariableDeclaratorSyntax;
                if ((decl.Parent as VariableDeclarationSyntax).Type.IsVar)
                {
                    return ExpressionContextKind.VAR_LOCAL_DECL;
                }
                Symbol sym = semModel.GetDeclaredSymbol(decl);
                if (sym is LocalSymbol)
                {
                    return ExpressionContextKind.LOCAL_DECL;
                }
                else if (sym is FieldSymbol)
                {
                    return ExpressionContextKind.FIELD_DECL;
                }
                else
                {
                    // ??? Is this reachable?
                    return ExpressionContextKind.UNKNOWN;
                }
            }
            else if (parent is FieldDeclarationSyntax)
            {
                return ExpressionContextKind.FIELD_DECL;
            }
            else if (parent is BinaryExpressionSyntax
                && (parent.Kind.Equals(SyntaxKind.LogicalAndExpression)
                || parent.Kind.Equals(SyntaxKind.LogicalOrExpression)))
            {
                return ExpressionContextKind.LOGICAL_ARG;
            }
            else if (parent is BinaryExpressionSyntax
                || parent.Kind.Equals(SyntaxKind.NegateExpression)
                /*&& (parent.Kind.Equals(SyntaxKind.PlusExpression)
                || parent.Kind.Equals(SyntaxKind.SubtractExpression)
                || parent.Kind.Equals(SyntaxKind.ModuloExpression)
                || parent.Kind.Equals(SyntaxKind.DivideExpression)
                || parent.Kind.Equals(SyntaxKind.MultiplyExpression)
                )*/)
            {
                return ExpressionContextKind.BINARY_ARG;
            }
            else if (parent.Kind.Equals(SyntaxKind.LogicalNotExpression))
            {
                return ExpressionContextKind.LOGICAL_ARG;
            }
            else if (parent is ParenthesizedLambdaExpressionSyntax || parent is SimpleLambdaExpressionSyntax)
            {
                return ExpressionContextKind.LAMBDA;
            }
            else if (parent is ReturnStatementSyntax || parent is YieldStatementSyntax)
            {
                return ExpressionContextKind.RETURN;
            }
            else if (parent is MemberAccessExpressionSyntax)
            {
                return ExpressionContextKind.LOOKUP;
            }
            else if (parent is InitializerExpressionSyntax)
            {
                return ExpressionContextKind.INITIALIZER;
            }
            else
            {
                // TODO
                return ExpressionContextKind.UNKNOWN;
            }
        }

        internal IEnumerable<int> AsEnumerable()
        {
            return allFeatures;
        }

        internal TypeSymbol ContainingType { get { return containingType; } }
    }
}
